"""
企业级金融数据MCP服务系统 - 数据库连接管理
提供PostgreSQL连接池、事务管理和ORM基础设施
"""

import asyncio
import logging
from typing import Optional, Dict, Any, List, AsyncGenerator, Union
from contextlib import asynccontextmanager
from datetime import datetime, timezone
import json

import asyncpg
from asyncpg import Pool, Connection, Record
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
from sqlalchemy import String, DateTime, Integer, Text, Boolean, JSON, Index
from sqlalchemy.dialects.postgresql import UUID
import uuid

from .config import config
from .exceptions import DatabaseConnectionError, DatabaseOperationError
from .logging import performance_logger, business_logger

logger = logging.getLogger(__name__)


class Base(DeclarativeBase):
    """SQLAlchemy基础模型类"""
    pass


class BaseModel(Base):
    """基础数据模型"""
    __abstract__ = True
    
    id: Mapped[str] = mapped_column(UUID(as_uuid=False), primary_key=True, default=lambda: str(uuid.uuid4()))
    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
    updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc), onupdate=lambda: datetime.now(timezone.utc))
    is_active: Mapped[bool] = mapped_column(Boolean, default=True)


class DataSyncLog(BaseModel):
    """数据同步日志模型"""
    __tablename__ = "data_sync_logs"
    
    data_source: Mapped[str] = mapped_column(String(100), nullable=False)
    sync_type: Mapped[str] = mapped_column(String(50), nullable=False)
    status: Mapped[str] = mapped_column(String(20), nullable=False)
    record_count: Mapped[Optional[int]] = mapped_column(Integer)
    error_message: Mapped[Optional[str]] = mapped_column(Text)
    sync_duration: Mapped[Optional[float]] = mapped_column()
    sync_metadata: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON)
    
    __table_args__ = (
        Index('idx_data_sync_logs_source_status', 'data_source', 'status'),
        Index('idx_data_sync_logs_created_at', 'created_at'),
    )


class ApiAccessLog(BaseModel):
    """API访问日志模型"""
    __tablename__ = "api_access_logs"
    
    endpoint: Mapped[str] = mapped_column(String(200), nullable=False)
    method: Mapped[str] = mapped_column(String(10), nullable=False)
    user_id: Mapped[Optional[str]] = mapped_column(String(100))
    ip_address: Mapped[Optional[str]] = mapped_column(String(45))
    user_agent: Mapped[Optional[str]] = mapped_column(Text)
    status_code: Mapped[Optional[int]] = mapped_column(Integer)
    response_time: Mapped[Optional[float]] = mapped_column()
    request_size: Mapped[Optional[int]] = mapped_column(Integer)
    response_size: Mapped[Optional[int]] = mapped_column(Integer)
    request_metadata: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON)
    
    __table_args__ = (
        Index('idx_api_access_logs_endpoint', 'endpoint'),
        Index('idx_api_access_logs_user_id', 'user_id'),
        Index('idx_api_access_logs_created_at', 'created_at'),
    )


class MCPToolExecutionLog(BaseModel):
    """MCP工具执行日志模型"""
    __tablename__ = "mcp_tool_execution_logs"
    
    tool_name: Mapped[str] = mapped_column(String(100), nullable=False)
    status: Mapped[str] = mapped_column(String(20), nullable=False)
    execution_time: Mapped[Optional[float]] = mapped_column()
    input_params: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON)
    output_size: Mapped[Optional[int]] = mapped_column(Integer)
    error_message: Mapped[Optional[str]] = mapped_column(Text)
    user_id: Mapped[Optional[str]] = mapped_column(String(100))
    session_id: Mapped[Optional[str]] = mapped_column(String(100))
    execution_metadata: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON)
    
    __table_args__ = (
        Index('idx_mcp_tool_logs_tool_status', 'tool_name', 'status'),
        Index('idx_mcp_tool_logs_user_id', 'user_id'),
        Index('idx_mcp_tool_logs_created_at', 'created_at'),
    )


class DatabaseManager:
    """数据库连接管理器"""
    
    def __init__(self):
        self._pool: Optional[Pool] = None
        self._engine = None
        self._session_factory = None
        self._initialized = False
    
    async def initialize(self):
        """初始化数据库连接"""
        if self._initialized:
            return
        
        try:
            # 创建asyncpg连接池
            self._pool = await asyncpg.create_pool(
                host=config.database.host,
                port=config.database.port,
                user=config.database.username,
                password=config.database.password,
                database=config.database.database,
                min_size=config.database.pool_size // 2,
                max_size=config.database.pool_size,
                command_timeout=config.database.command_timeout,
                server_settings={
                    'application_name': 'financial_mcp_service',
                    'timezone': 'UTC'
                }
            )
            
            # 创建SQLAlchemy异步引擎
            database_url = f"postgresql+asyncpg://{config.database.username}:{config.database.password}@{config.database.host}:{config.database.port}/{config.database.database}"
            self._engine = create_async_engine(
                database_url,
                pool_size=config.database.pool_size,
                max_overflow=config.database.max_overflow,
                pool_timeout=config.database.pool_timeout,
                pool_recycle=config.database.pool_recycle,
                echo=config.database.echo,
                pool_pre_ping=config.database.pool_pre_ping
            )
            
            # 创建会话工厂
            self._session_factory = async_sessionmaker(
                self._engine,
                class_=AsyncSession,
                expire_on_commit=False
            )
            
            # 测试连接
            await self._test_connection()
            
            self._initialized = True
            logger.info("数据库连接初始化成功")
            
        except Exception as e:
            logger.error(f"数据库连接初始化失败: {e}")
            raise DatabaseConnectionError(f"数据库连接初始化失败: {e}")
    
    async def _test_connection(self):
        """测试数据库连接"""
        if not self._pool:
            raise DatabaseConnectionError("连接池未初始化")
        
        async with self._pool.acquire() as conn:
            result = await conn.fetchval("SELECT 1")
            if result != 1:
                raise DatabaseConnectionError("数据库连接测试失败")
    
    async def close(self):
        """关闭数据库连接"""
        if self._pool:
            await self._pool.close()
            self._pool = None
        
        if self._engine:
            await self._engine.dispose()
            self._engine = None
        
        self._session_factory = None
        self._initialized = False
        logger.info("数据库连接已关闭")
    
    @asynccontextmanager
    async def get_connection(self) -> AsyncGenerator[Connection, None]:
        """获取数据库连接"""
        if not self._pool:
            raise DatabaseConnectionError("数据库连接池未初始化")
        
        async with self._pool.acquire() as conn:
            try:
                yield conn
            except Exception as e:
                logger.error(f"数据库操作错误: {e}")
                raise DatabaseOperationError(f"数据库操作错误: {e}")
    
    @asynccontextmanager
    async def get_session(self) -> AsyncGenerator[AsyncSession, None]:
        """获取SQLAlchemy会话"""
        if not self._session_factory:
            raise DatabaseConnectionError("数据库会话工厂未初始化")
        
        async with self._session_factory() as session:
            try:
                yield session
            except Exception as e:
                await session.rollback()
                logger.error(f"数据库会话错误: {e}")
                raise DatabaseOperationError(f"数据库会话错误: {e}")
    
    @asynccontextmanager
    async def transaction(self) -> AsyncGenerator[Connection, None]:
        """事务上下文管理器"""
        async with self.get_connection() as conn:
            async with conn.transaction():
                yield conn
    
    async def execute_query(self, 
                           query: str, 
                           params: Optional[Union[Dict[str, Any], List[Any]]] = None,
                           fetch_mode: str = "all") -> Union[List[Record], Record, Any]:
        """执行查询"""
        with performance_logger.timer(f"database_query"):
            async with self.get_connection() as conn:
                if fetch_mode == "all":
                    return await conn.fetch(query, *(params or []))
                elif fetch_mode == "one":
                    return await conn.fetchrow(query, *(params or []))
                elif fetch_mode == "val":
                    return await conn.fetchval(query, *(params or []))
                else:
                    raise ValueError(f"不支持的fetch_mode: {fetch_mode}")
    
    async def execute_command(self, 
                             command: str, 
                             params: Optional[Union[Dict[str, Any], List[Any]]] = None) -> str:
        """执行命令（INSERT, UPDATE, DELETE等）"""
        with performance_logger.timer(f"database_command"):
            async with self.get_connection() as conn:
                return await conn.execute(command, *(params or []))
    
    async def bulk_insert(self, 
                         table_name: str, 
                         records: List[Dict[str, Any]], 
                         conflict_resolution: str = "ignore") -> int:
        """批量插入数据"""
        if not records:
            return 0
        
        with performance_logger.timer(f"database_bulk_insert_{table_name}"):
            async with self.get_connection() as conn:
                columns = list(records[0].keys())
                values = [[record[col] for col in columns] for record in records]
                
                if conflict_resolution == "ignore":
                    query = f"""
                        INSERT INTO {table_name} ({', '.join(columns)})
                        SELECT * FROM unnest($1::{', '.join(['text'] * len(columns))})
                        ON CONFLICT DO NOTHING
                    """
                elif conflict_resolution == "update":
                    update_clause = ', '.join([f"{col} = EXCLUDED.{col}" for col in columns if col != 'id'])
                    query = f"""
                        INSERT INTO {table_name} ({', '.join(columns)})
                        SELECT * FROM unnest($1::{', '.join(['text'] * len(columns))})
                        ON CONFLICT (id) DO UPDATE SET {update_clause}
                    """
                else:
                    query = f"""
                        INSERT INTO {table_name} ({', '.join(columns)})
                        SELECT * FROM unnest($1::{', '.join(['text'] * len(columns))})
                    """
                
                result = await conn.execute(query, values)
                return int(result.split()[-1])
    
    async def get_table_stats(self, table_name: str) -> Dict[str, Any]:
        """获取表统计信息"""
        query = """
            SELECT 
                schemaname,
                tablename,
                attname,
                n_distinct,
                correlation
            FROM pg_stats 
            WHERE tablename = $1
        """
        
        stats = await self.execute_query(query, [table_name])
        
        count_query = f"SELECT COUNT(*) FROM {table_name}"
        total_count = await self.execute_query(count_query, fetch_mode="val")
        
        return {
            "table_name": table_name,
            "total_rows": total_count,
            "column_stats": [dict(stat) for stat in stats]
        }
    
    async def create_tables(self):
        """创建数据表"""
        if not self._engine:
            raise DatabaseConnectionError("数据库引擎未初始化")
        
        async with self._engine.begin() as conn:
            await conn.run_sync(Base.metadata.create_all)
        
        logger.info("数据表创建完成")


class QueryBuilder:
    """SQL查询构建器"""
    
    def __init__(self, table_name: str):
        self.table_name = table_name
        self._select_fields = ["*"]
        self._where_conditions = []
        self._order_by = []
        self._limit_value = None
        self._offset_value = None
        self._params = []
    
    def select(self, *fields: str) -> "QueryBuilder":
        """选择字段"""
        self._select_fields = list(fields)
        return self
    
    def where(self, condition: str, *params: Any) -> "QueryBuilder":
        """添加WHERE条件"""
        self._where_conditions.append(condition)
        self._params.extend(params)
        return self
    
    def order_by(self, field: str, direction: str = "ASC") -> "QueryBuilder":
        """添加排序"""
        self._order_by.append(f"{field} {direction}")
        return self
    
    def limit(self, count: int) -> "QueryBuilder":
        """设置LIMIT"""
        self._limit_value = count
        return self
    
    def offset(self, count: int) -> "QueryBuilder":
        """设置OFFSET"""
        self._offset_value = count
        return self
    
    def build(self) -> tuple[str, List[Any]]:
        """构建查询"""
        query_parts = [f"SELECT {', '.join(self._select_fields)} FROM {self.table_name}"]
        
        if self._where_conditions:
            query_parts.append(f"WHERE {' AND '.join(self._where_conditions)}")
        
        if self._order_by:
            query_parts.append(f"ORDER BY {', '.join(self._order_by)}")
        
        if self._limit_value:
            query_parts.append(f"LIMIT {self._limit_value}")
        
        if self._offset_value:
            query_parts.append(f"OFFSET {self._offset_value}")
        
        return " ".join(query_parts), self._params


# 全局数据库管理器实例
db_manager = DatabaseManager()


async def init_database():
    """初始化数据库"""
    await db_manager.initialize()
    await db_manager.create_tables()
    business_logger.log_system_event("database_init", "数据库初始化完成")


async def close_database():
    """关闭数据库连接"""
    await db_manager.close()
    business_logger.log_system_event("database_close", "数据库连接已关闭")